# -*- coding: UTF-8 -*-
import os
import netCDF4 as nc
import numpy as np
from osgeo import gdal, osr, ogr
import glob
from tqdm import tqdm

# os.environ['PROJ_LIB'] = r'./proj'

# 用于逐月数据集

def get_geotransform(nc_file: str, key:str):
    data = nc.Dataset(nc_file)  # 利用.Dataset()读取nc数据
    Lat = data.variables['lat'][:]
    Lon = data.variables['lon'][:]
    data_arr = np.asarray(data.variables[key])  # 属性变量名
    # 影像的左上角&右下角坐标
    Lonmin, Latmax, Lonmax, Latmin = [Lon.min(), Lat.max(), Lon.max(), Lat.min()]
    # Lonmin, Latmax, Lonmax, Latmin
    # 分辨率计算
    Num_lat = len(Lat)
    Num_lon = len(Lon)
    Lat_res = (Latmax - Latmin) / (float(Num_lat) - 1)
    Lon_res = (Lonmax - Lonmin) / (float(Num_lon) - 1)
    # 设置影像的显示范围
    # Lat_re前需要添加负号
    nodata =data.variables[key].missing_value  # netcdf文件的缺失值
    geotransform = (Lonmin, Lon_res, 0.0, Latmax, 0.0, -Lat_res)
    mask = np.where(data_arr[0,:,:] == nodata, 0, 1)  # nodata区域为 0, 多波段只读取一个波段
    return geotransform, nodata, mask

def arr2tif(data_arr, out_tif_name: str, geotransform,nodata):
    # 创建tif文件
    driver = gdal.GetDriverByName('GTiff')
    out_tif = driver.Create(out_tif_name, data_arr.shape[1],data_arr.shape[0],1, gdal.GDT_Float32)
    # driver.Create(filename, width, height, bands, DataType.GDT_Byte, null)
    out_tif.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    out_tif.SetProjection(prj.ExportToWkt())
    out_tif.GetRasterBand(1).SetNoDataValue(nodata)
    # 数据导出
    out_tif.GetRasterBand(1).WriteArray(data_arr)  # 将数据写入内存
    out_tif.FlushCache()  # 将数据写入到硬盘
    out_tif = None  # 关闭tif文件

    return out_tif_name

# #测试数据信息，每一个nc文件是逐月降水、均温数据
# #中国1km分辨率逐月降水量数据集（1901 - 2023）
# #https://data.tpdc.ac.cn/zh-hans/data/faae7605-a0f2-4d18-b28f-5cee413766a2
#中国1km分辨率逐月平均气温数据集（1901-2023）
#https://data.tpdc.ac.cn/zh-hans/data/71ab4677-b66c-4fd1-a004-b2a541c4d5bf
if __name__ == "__main__":

    ##++++++++++++++++++++++++循环处理单张图像+++++++++++++++++++++++++++++++++++
    # Input_nc = r"F:\test\MRI-ESM2-0\MRI-ESM2-0_ssp585_tmp-30s-20.nc"
    # Output_tif = r"F:\test\tmp_2060_ssp585.tif"
    # key1 = 'etp' #'tmp''pre'
    # band_dimension = 24  # 时间维度，12个月
    # scale = 0.1  # 降水比例转换0.1mm，温度比例转换 0.1deg，pet单位0.1mm
    #
    # geotran, nodata, mask= get_geotransform(Input_nc, key1)
    # _mask = np.where(mask == 0, 1, 0)  # mask的反
    # pre_arr = np.asarray(nc.Dataset(Input_nc).variables[key1]) * scale
    # if band_dimension == 24:
    #     pre_arr = pre_arr[12:]  # 删除前12个月，因为前12个月是上一年的数据
    # M = pre_arr.sum(axis=0) * mask + _mask * nodata  # 累积年值，把nodata数据添加上
    # arr2tif(M, Output_tif, geotran, nodata)  # 保存年值为tif
    # print("计算完成！")

    # ++++++++++++++++++++++++循环处理多幅图像+++++++++++++++++++++++++++++++++++
    Input_folder = r"D:\数据处理\TMP"
    Output_folder = r"D:\数据处理\TMP_a"
    start_year = 2000
    end_year = 2023
    key1 = 'tmp'  # 均温是，降水是pre，pet是etp
    band_dimension = 12  # 时间维度，12个月
    scale = 0.1  # 降水比例转换0.1mm，温度比例转换 0.1deg，pet单位0.1mm

    if not os.path.exists(Output_folder):  # 如果目录不存在就返回False，创建该文件夹
        os.mkdir(Output_folder)

    # 读取所有数据
    nc_list = sorted(glob.glob(os.path.join(Input_folder, f"*.nc")))

    geotran, nodata, mask= get_geotransform(nc_list[1], key1)
    _mask = np.where(mask == 0, 1, 0)  # mask的反
    # 建一个循环，每个循环分别读取月降水和月均温，计算逐月Rh，最后sum计算年Rh
    if band_dimension == 12:
        for k in tqdm(range(start_year, end_year+1)):  # 循环每一个nc文件，内有12个波段/dims
            j = k-start_year
            pre_arr = np.asarray(nc.Dataset(nc_list[j]).variables[key1]) * scale
            M = pre_arr.sum(axis=0) * mask + _mask * nodata  # 累积年值，把nodata数据添加上
            if key1 == 'tmp': # 均温需要除以12
                M=M/12.0
            arr2tif(M, os.path.join(Output_folder, key1 + "_"f"{k}.tif"), geotran, nodata)  # 保存年值为tif
            print(f"{k}年"+key1+"计算完成！")
    elif band_dimension == 24:
        for k in tqdm(range(start_year, end_year + 1, 2)):  # 循环每一个nc文件，内有12个波段/dims
            j = k - start_year
            # 读取前一年的前12个波段
            pre_arr_first_half = np.asarray(nc.Dataset(nc_list[j]).variables[key1])[:12] * scale
            # 读取后一年的后12个波段
            pre_arr_second_half = np.asarray(nc.Dataset(nc_list[j + 1]).variables[key1])[12:] * scale

            # 计算两年和值
            M1 = pre_arr_first_half.sum(axis=0) * mask + _mask * nodata  # 累积年值，把nodata数据添加上
            M2 = pre_arr_second_half.sum(axis=0) * mask + _mask * nodata

            if key1 == 'tmp': # 均温需要除以12
                M1=M2/12.0
                M2=M2/12.0
            arr2tif(M1, os.path.join(Output_folder, key1 + "_"f"{k}.tif"), geotran, nodata)  # 保存年值为tif
            print(f"{k}年" + key1 + "计算完成！")
            arr2tif(M2, os.path.join(Output_folder, key1 + "_"f"{k+1}.tif"), geotran, nodata)
            print(f"{k+1}年"+key1+"计算完成！")

